import os
import dashscope
from dashscope import Generation
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 设置API Key
dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")

def qwen_chat(prompt, tools=None):
    """调用通义千问模型进行对话"""
    response = Generation.call(
        model="qwen-max",
        prompt=prompt,
        tools=tools
    )
    return response.output.text

def get_weather(city):
    """模拟天气查询工具"""
    # 这里应该是调用真实天气API的代码
    # 为了简化，我们模拟返回天气信息
    return f"{city}今天晴天，25度，微风"

def main():
    print("欢迎使用AI问答助手！输入'退出'结束对话。")
    
    # 定义可用的工具
    tools = [
        {
            "name": "get_weather",
            "description": "获取指定城市的天气信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "城市名称"
                    }
                },
                "required": ["city"]
            }
        }
    ]
    
    while True:
        user_input = input("\n你: ")
        if user_input.lower() == "退出":
            print("再见！")
            break
            
        # 检查是否需要调用工具
        if "天气" in user_input:
            # 提取城市名（简化处理）
            city = user_input.replace("天气", "").strip()
            if not city:
                city = "北京"  # 默认城市
                
            # 调用工具
            weather_info = get_weather(city)
            
            # 将工具结果提供给模型
            response = qwen_chat(f"用户问: {user_input}\n工具返回: {weather_info}")
        else:
            # 直接调用模型
            response = qwen_chat(user_input)
            
        print(f"AI: {response}")

if __name__ == "__main__":
    main()
